/**************************************************************************************************
	$Header: /pub/cvsroot/mydns/src/test/main.c,v 1.13 2002/07/25 21:43:44 bboy Exp $

	Copyright (C) 2002  Don Moore <bboy@bboy.net>

	This program is free software; you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation; either version 2 of the License, or
	(at Your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program; if not, write to the Free Software
	Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
**************************************************************************************************/

#include "mydnsutil.h"
#include "../mydns/bits.h"
#include "../mydns/header.h"
#include "mydns.h"
#include <netdb.h>


/* Config */
char	inputfile[PATH_MAX] = "TESTDATA";					/* Name of input file */
char	host[256] = "localhost";								/* Hostname of DNS server */
int	port = 55;													/* Port number of DNS server */
int	timeout = 1;												/* Time out seconds */
int	garbage = 0;												/* Run garbage test? */
int	maxreq = 0;													/* Number of questions to ask, 0=Inf */

/* Totals */
int	req;															/* Number of requests */

/* Misc vars */
int	lineno;														/* Current line number */
FILE	*fp;															/* Input file pointer */
struct sockaddr_in sa;											/* Address of remote server */
struct timeval tvs;												/* Time testing started */
int timed_out;

extern void send_garbage(void);


/*+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
	USAGE
	Display program usage information.
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++*/
static void
usage(status)
	int status;
{
	if (status != EXIT_SUCCESS)
	{
		fprintf(stderr, _("Try `%s --help' for more information."), progname);
		fputs("\n", stderr);
	}
	else
	{
		printf(_("Usage: %s [HOST:PORT] [INPUTFILE]"), progname);
		puts("");
		puts(_("DNS server testing/benchmarking utility."));
		puts("");
/*		puts("----------------------------------------------------------------------------78");  */
		puts(_("  -r, --random   start from a random position in INPUTFILE"));
		puts(_("  -g, --garbage  ignore INPUTFILE, run garbage test instead"));
		puts(_("  -n, --num=NUM  ask NUM questions then exit"));
		puts(_("      --help     display this help and exit"));
		puts(_("      --version  output version information and exit"));
		puts("");
		puts(_("This program reads INPUTFILE which must be a comma-delimited file"));
		puts(_("containing any number of lines with the format:"));
		puts("");
		puts(_("   QUERYTYPE,name    (Ex: \"MX,example.com\")"));
		puts("");
		puts(_("Each query in the file will be submitted to the DNS server at HOST:PORT"));
		puts(_("as quickly as possible."));
		puts("");
		printf(_("The default input file is \"%s\".\n"), inputfile);
		printf(_("The default DNS server is %s:%u.\n"), host, port);
		puts("");
		printf(_("Report bugs to <%s>.\n"), PACKAGE_BUGREPORT);
	}
	exit(status);
}
/*--- usage() -----------------------------------------------------------------------------------*/


/*+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
	CMDLINE
	Process command line options.
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++*/
static void
cmdline(argc, argv)
	int argc;
	char *argv[];
{
	char	*optstr;
	int	optc, optindex, jumprandom = 0;
	struct stat st;
	struct hostent *he;
	struct option const longopts[] =
	{
		{"num",				required_argument,	NULL,	'n'},
		{"random",			no_argument,			NULL,	'r'},
		{"garbage",			no_argument,			NULL,	'g'},
		{"help",				no_argument,			NULL,	0},
		{"version",			no_argument,			NULL,	0},

		{NULL, 0, NULL, 0}
	};

	srand(time(NULL) + getpid());
	error_init(argv[0], LOG_USER);							/* Init output routines */

	optstr = getoptstr(longopts);
	while ((optc = getopt_long(argc, argv, optstr, longopts, &optindex)) != -1)
	{
		switch (optc)
		{
			case 0:
				{
					const char *opt = longopts[optindex].name;

					if (!strcmp(opt, "version"))									/* --version */
					{
						printf("%s ("PACKAGE_NAME") "PACKAGE_VERSION"\n", progname);
						puts("\n" PACKAGE_COPYRIGHT);
						puts(_("This is free software; see the source for copying conditions.  There is NO"));
						puts(_("warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE."));
						exit(EXIT_SUCCESS);
					}
					else if (!strcmp(opt, "help"))								/* --help */
						usage(EXIT_SUCCESS);
				}
				break;
			case 'g':
				garbage = 1;
				break;
			case 'n':
				maxreq = atoi(optarg);
				break;
			case 'r':
				jumprandom = 1;
				break;
			default:
				usage(EXIT_FAILURE);
		}
	}
	if (optind < argc)		/* First arg: Hostname:Port */
	{
		char *p;

		strncpy(host, argv[optind], sizeof(host)-1);
		if ((p = strrchr(host, ':')))
		{
			*p++ = '\0';
			port = atoi(p);
		}
		optind++;
	}
	if (optind < argc)		/* Second arg: Filename */
	{
		strncpy(inputfile, argv[optind], sizeof(inputfile)-1);
		optind++;
	}
	if (optind < argc)
		Warnx("%s", _("Extraneous command-line argument(s) ignored"));

	if (!garbage)
	{
		/* Open input file */
		if (inputfile[0] == '-' && !inputfile[1])
			fp = stdin;
		else
			if (!(fp = fopen(inputfile, "r")))
				Err("%s", inputfile);

		/* Jump to random pos if requested */
		if (jumprandom)
		{
			char buf[BUFSIZ];
			off_t off;

			stat(inputfile, &st);
			off = rand() % (int)(st.st_size * 0.75);
			fseek(fp, off, SEEK_SET);
			fgets(buf, sizeof(buf), fp);
			Notice("%s: %s %lu/%lu (%.0f%%)", inputfile, _("starting offset"),
					 off, st.st_size, PCT(st.st_size, off));
		}
	}

	/* Get server address */
	memset(&sa, 0, sizeof(struct sockaddr_in));
	sa.sin_family = AF_INET;
	sa.sin_port = htons(port);
	if (!(he = gethostbyname(host)))
		Errx("%s: %s", host, _("unknown host"));
	memcpy(&sa.sin_addr, he->h_addr_list[0], sizeof(struct in_addr));
}
/*--- cmdline() ---------------------------------------------------------------------------------*/


/*+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
	CURRENT_RATE
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++*/
static double
current_rate(void)
{
	struct timeval tv;

	gettimeofday(&tv, NULL);
	return (req / ((tv.tv_sec + tv.tv_usec / 1000000.0) - (tvs.tv_sec + tvs.tv_usec / 1000000.0)));
}
/*--- current_rate() ----------------------------------------------------------------------------*/


/*+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
	SUMMARY_EXIT
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++*/
void
summary_exit(int signo)
{
	struct timeval tv;
	double secs, rs;

	puts("");
	if (signo == SIGALRM)
	{
		printf(_("timed out\n"));
		timed_out = 1;
		alarm(timeout);
		return;
	}
	gettimeofday(&tv, NULL);
	secs = ((tv.tv_sec + tv.tv_usec / 1000000.0) - (tvs.tv_sec + tvs.tv_usec / 1000000.0));
	rs = req / secs;
	Notice(_("%s:%u: %u questions in %.2f seconds (%.2f/s)"), host, port, req, secs, rs);
	exit(0);
}
/*--- summary_exit() ----------------------------------------------------------------------------*/


/*+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
	MAKE_QUESTION
	Creates a question.  Returns the packet and stores the length of the packet in `packetlen'.
	The packet is dynamically allocated and should be free()'d.
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++*/
u_char *
make_question(u_int16_t id, dns_qtype_t qtype, u_char *name, size_t *packetlen)
{
	static u_char req[1024];
	u_char *dest = req, *c;
	DNS_HEADER	header;
	size_t len;

	if (packetlen) *packetlen = 0;

	memset(&header, 0, sizeof(DNS_HEADER));
	DNS_PUT16(dest, id);											/* ID */
	header.rd = 1;
	memcpy(dest, &header, sizeof(DNS_HEADER)); dest += SIZE16;
	DNS_PUT16(dest, 1);											/* QDCOUNT */
	DNS_PUT16(dest, 0);											/* ANCOUNT */
	DNS_PUT16(dest, 0);											/* NSCOUNT */
	DNS_PUT16(dest, 0);											/* ARCOUNT */
	for (c = name; *c; c++)										/* QNAME */
		if (c == name || *c == '.')
		{
			u_char *end;
			if (c != name)
				c++;
			if ((end = strchr(c, '.')))
				*end = '\0';
			if ((len = strlen(c)))
			{
				if (len > 64)
				{
					Warnx(_("zone contains invalid label (64 chars max)"));
					return (NULL);
				}
				*dest++ = len;
				DNS_PUT(dest, c, len);
			}
			if (end)
				*end = '.';
		}
	*dest++ = 0;
	DNS_PUT16(dest, (u_int16_t)qtype);						/* QTYPE */
	DNS_PUT16(dest, DNS_C_IN);									/* QCLASS */
	len = dest - req;

	if (packetlen) *packetlen = len;

	return (req);
}
/*--- make_question() ---------------------------------------------------------------------------*/


/*+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
	SEND_DNS_REQUEST
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++*/
static void
send_dns_request(const char *type, const char *name)
{
	static int fd = -1;
	int rv, addrlen;
	u_char buf[BUFSIZ], *b;
	u_char *question;
	size_t len;
	dns_qtype_t qtype = DNS_T_NONE;
	static u_int16_t id = 1;

	if (fd != -1)
		close(fd);
	timed_out = 0;
	alarm(timeout);

	/* Connect to server */
	if ((fd = socket(AF_INET, SOCK_DGRAM, 0)) < 0)
		Err("%s:%u", host, port);
	if (timed_out)
		return;
	if (connect(fd, (const struct sockaddr *)&sa, sizeof(struct sockaddr_in)))
		Err("%s:%u", host, port);
	if (timed_out)
		return;

	/* Flush any pending input */
	do
	{
		fd_set rfd;
		struct timeval tv;

		FD_ZERO(&rfd);
		FD_SET(fd, &rfd);
		tv.tv_sec = tv.tv_usec = 0;
		if ((rv = select(fd+1, &rfd, NULL, NULL, &tv)) > 0)
			recvfrom(fd, &buf, sizeof(buf), 0, (struct sockaddr *)&sa, &addrlen);
	} while (rv > 0);

	strncpy(buf, name, sizeof(buf)-1);
	if (LASTCHAR(buf) != '.')
		strcat(buf, ".");

	if (++id > 50000)
		id = 1;

	if (!strcasecmp(type, "A")) qtype = DNS_T_A;
	else if (!strcasecmp(type, "ANY")) qtype = DNS_T_ANY;
	else if (!strcasecmp(type, "MX")) qtype = DNS_T_MX;
	else if (!strcasecmp(type, "SOA")) qtype = DNS_T_SOA;
	else if (!strcasecmp(type, "AAAA")) qtype = DNS_T_AAAA;
	else if (!strcasecmp(type, "CNAME")) qtype = DNS_T_CNAME;
	else if (!strcasecmp(type, "PTR")) qtype = DNS_T_PTR;
	else if (!strcasecmp(type, "NS")) qtype = DNS_T_NS;
	else if (!strcasecmp(type, "TXT")) qtype = DNS_T_TXT;
	else return (void)Warnx("%s:%u: %s: %s", inputfile, lineno, type, _("unsupported query type"));

	question = make_question(id, qtype, (u_char *)name, &len);
	if ((rv = sendto(fd, question, len, 0, (struct sockaddr *)&sa, sizeof(sa))) < 0)
	{
		Warn("%s:%u: %s:%u: %s", inputfile, lineno, host, port, _("error sending question"));
		return;
	}
	if (timed_out)
		return;
	if (!rv)
		return;

	/* Read reply */
	addrlen = sizeof(sa);
	if ((rv = recvfrom(fd, &buf, sizeof(buf), 0, (struct sockaddr *)&sa, &addrlen)) == -1)
		Err("%s:%u: %s:%u: %s", inputfile, lineno, host, port, _("error reading reply"));
	req++;
	if (timed_out)
		return;

	if (!rv)
		fprintf(stderr, "%s:%u: %s %s: %s\n", inputfile, lineno, type, name, _("no reply"));

	b = buf;
	{
		u_int16_t rcv_id;
		DNS_HEADER hdr;

		DNS_GET16(rcv_id, b);
		if (rcv_id != id)
			fprintf(stderr, "%s:%u: %s:%s: %s (got=%u real=%u)\n", inputfile, lineno, type, name, _("id mismatch"), rcv_id, id);
		memcpy(&hdr, b, SIZE16); b += SIZE16;
		if (hdr.rcode != DNS_R_OK)
			fprintf(stderr, "%s:%u: %s:%s: %s\n", inputfile, lineno, type, name, mydns_rcode_str(hdr.rcode));
	}
}
/*--- send_dns_request() ------------------------------------------------------------------------*/


/*+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
	MAIN
+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++*/
int
main(argc, argv)
	int argc;
	const char *argv[];
{
	char buf[BUFSIZ], *d;

	bindtextdomain(PACKAGE, LOCALEDIR);						/* Internationalization */
	cmdline(argc, argv);											/* Process command line */
	signal(SIGINT, summary_exit);
	signal(SIGALRM, summary_exit);

	gettimeofday(&tvs, NULL);

	if (garbage)
		send_garbage();

	while (fgets(buf, sizeof(buf), fp))
	{
		if (lineno && !(lineno % 4000))
			printf("%7d %.1f\n", lineno, current_rate());
		else if (lineno && !(lineno % 100))
			printf("."), fflush(stdout);
		lineno++;

		if (!(d = strchr(buf, ',')))
		{
			Warnx("%s:%u: %s", inputfile, lineno, _("input line does not contain a comma"));
			continue;
		}
		strtrim(buf);
		*d++ = '\0';
		send_dns_request(buf, d);

		if (req == maxreq)
			break;
	}
	summary_exit(0);
	return (0);
}
/*--- main() ------------------------------------------------------------------------------------*/

/* vi:set ts=3: */
/* NEED_PO */
